/***************************************************************************
 *   Copyright (C) 2011-2012 by Francesco Cecconi                          *
 *   francesco.cecconi@gmail.com                                           *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License.        *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
 ***************************************************************************/

#include "vulnerability.h"
#include "mainwindow.h"

WelcomeVulnerability::WelcomeVulnerability(QWidget* parent) : QWidget(parent)
{
    setupUi(this);
}

VulnerabilityWidget::VulnerabilityWidget(QWidget* parent): QWidget(parent)
{
    setupUi(this);
}

Vulnerability::Vulnerability(MainWindow* parent)
    : QObject(parent), m_ui(parent)
{

    m_vulnerabilityWidget = new VulnerabilityWidget(m_ui);

    m_vulnerabilityWidget->tWresult->setTabsClosable(true);
    m_vulnerabilityWidget->tWresult->removeTab(0);
    m_vulnerabilityWidget->tWresult->setVisible(false);
    updateCompleter();
    loadSearchUrlIntoCombo();

    m_vulnerabilityWidget->comboVulnRis->setStyleSheet(QString::fromUtf8("color: rgb(153, 153, 153);"));
    m_vulnerabilityWidget->comboVulnRis->insertItem(0, tr("Search Vulnerabilities"));

    setupActions();

    connect(m_vulnerabilityWidget->tWresult, SIGNAL(tabCloseRequested(int)),
            this, SLOT(closeVulnTab(int)));
    connect(m_vulnerabilityWidget->comboVulnRis->lineEdit(), SIGNAL(returnPressed()),
            this, SLOT(searchVulnerabilityFromCombo()));
    connect(m_ui->actionAdd_vulnerability_search_url, SIGNAL(triggered()),
            this, SLOT(showAddUrlUi()));
    connect(m_vulnerabilityWidget->comboVulnRis->lineEdit(), SIGNAL(cursorPositionChanged(int,int)),
            this, SLOT(updateComboServicesProperties()));
    connect(m_vulnerabilityWidget->comboVuln, SIGNAL(currentIndexChanged(QString)),
            this, SLOT(updateComboVuln(QString)));
    connect(m_vulnerabilityWidget->treeVulnNseRecovered, SIGNAL(itemClicked(QTreeWidgetItem*, int)),
            this, SLOT(searchVulnFromTreeWidget()));

    m_mainHorizontalSplitter = new QSplitter(m_ui);

    m_mainHorizontalSplitter->setOrientation(Qt::Horizontal);
    m_mainHorizontalSplitter->addWidget(m_vulnerabilityWidget->treeVulnNseRecovered);
    m_mainHorizontalSplitter->addWidget(m_vulnerabilityWidget->frameVuln);

    m_vulnerabilityWidget->layout()->addWidget(m_mainHorizontalSplitter);

    QSettings settings("nmapsi4", "nmapsi4");
    if (!settings.value("vulnTreeSplitterHorizontal").toByteArray().isEmpty()) {
        m_mainHorizontalSplitter->restoreState(settings.value("vulnTreeSplitterHorizontal").toByteArray());
    } else {
        // set a default width
        QList<int> size = m_mainHorizontalSplitter->sizes();
        size[0] = 100; // ratio
        size[1] = 200;
        m_mainHorizontalSplitter->setSizes(size);
    }

    m_welcomeUi = new WelcomeVulnerability(m_ui);
    m_vulnerabilityWidget->frameVuln->layout()->addWidget(m_welcomeUi);
}

Vulnerability::~Vulnerability()
{
    freemap<QString, QAction*>::itemDeleteAll(m_collections);
}

void Vulnerability::syncSettings()
{
    QSettings settings("nmapsi4", "nmapsi4");
    settings.setValue("vulnTreeSplitterHorizontal", m_mainHorizontalSplitter->saveState());
}

void Vulnerability::updateCompleter()
{
    // load vulnerability services string model
    if (!m_ui->m_bookmark->isBookmarkServicesListEmpty()) {
        if (!m_completerVuln.isNull()) {
            QStringListModel *newModel = qobject_cast<QStringListModel*>(m_completerVuln.data()->model());
            newModel->setStringList(m_ui->m_bookmark->getServicesListFromBookmark());
        } else if (m_vulnModel.isNull()) {
            m_vulnModel = new QStringListModel(m_ui->m_bookmark->getServicesListFromBookmark(), this);
        }
    }
}

void Vulnerability::setupActions()
{
    /**
     * Setup actions for vulnerability section
     */
    QAction* action;

    action = new QAction(this);
    action->setIcon(QIcon(QString::fromUtf8(":/images/images/viewmag.png")));
    action->setIconText(tr("Search"));
    action->setEnabled(false);
    m_collections.insert("search-act", action);
    connect(action, SIGNAL(triggered()), this, SLOT(searchVulnerabilityFromCombo()));
    m_ui->toolBarSearch->addAction(action);

    action = new QAction(this);
    action->setIcon(QIcon(QString::fromUtf8(":/images/images/go-previous.png")));
    action->setIconText(tr("Back"));
    action->setEnabled(false);
    m_collections.insert("back-act", action);
    connect(action, SIGNAL(triggered()), this, SLOT(tabWebBack()));
    m_ui->toolBarSearch->addAction(action);

    action = new QAction(this);
    action->setIcon(QIcon(QString::fromUtf8(":/images/images/go-next.png")));
    action->setIconText(tr("Forward"));
    action->setEnabled(false);
    m_collections.insert("forward-act", action);
    connect(action, SIGNAL(triggered()), this, SLOT(tabWebForward()));
    m_ui->toolBarSearch->addAction(action);

    action = new QAction(this);
    action->setIcon(QIcon(QString::fromUtf8(":/images/images/button_cancel.png")));
    action->setIconText(tr("Stop"));
    action->setEnabled(false);
    m_collections.insert("stop-act", action);
    connect(action, SIGNAL(triggered()), this, SLOT(tabWebStop()));
    m_ui->toolBarSearch->addAction(action);

    m_ui->toolBarSearch->setContextMenuPolicy(Qt::PreventContextMenu);
}


const QList< QPair<QString, QString> > Vulnerability::getDefaultUrlList()
{
    QList< QPair<QString, QString> > urlListModel;
    QPair<QString, QString> urlModel;

    // set default name+value url
    urlModel.first = "cve.mitre.org";
    urlModel.second = "http://cve.mitre.org/cgi-bin/cvekey.cgi?keyword=";
    urlListModel.push_back(urlModel);

    urlModel.first = "www.securityfocus.com";
    urlModel.second = "http://www.securityfocus.com/swsearch?sbm=bid&submit=Search%21&metaname=alldoc&sort=swishrank&query=";
    urlListModel.push_back(urlModel);

    urlModel.first = "secunia.com";
    urlModel.second = "http://secunia.com/advisories/search/?search=";
    urlListModel.push_back(urlModel);

    urlModel.first = "web.nvd.nist.gov";
    urlModel.second = "http://web.nvd.nist.gov/view/vuln/search-results?search_type=all&cves=on&query=";
    urlListModel.push_back(urlModel);

    urlModel.first = "osvdb.org";
    urlModel.second = "http://osvdb.org/search?search[vuln_title]=";
    urlListModel.push_back(urlModel);

    return urlListModel;
}

void Vulnerability::searchVulnerabilityFromCombo()
{
    if (m_vulnerabilityWidget->comboVulnRis->currentText().isEmpty()) {
        return;
    }

    // save current search on bookmark
    m_ui->m_bookmark->saveServiceToBookmark(m_vulnerabilityWidget->comboVulnRis->currentText(), m_ui->m_hostCache);
    //update completer cache
    updateCompleter();

    QString finalUrlString;
    QString tmpSearchString;
    // search string from QComboBox
    tmpSearchString = m_vulnerabilityWidget->comboVulnRis->currentText();
    // space to + for correct search string format
    tmpSearchString.replace(QString(" "), QString("+"));

    // create QHash for default profile
    QList< QPair<QString, QString> > urlListModel = getDefaultUrlList();

    // selection url from key for search
    if (m_vulnerabilityWidget->comboWebV->currentIndex() < urlListModel.size()) {
        // It is a default url, stored in QList
        QListIterator< QPair<QString, QString> > i(urlListModel);
        while (i.hasNext()) {
            QPair<QString, QString> urlModel = i.next();
            if (urlModel.first == m_vulnerabilityWidget->comboWebV->currentText()) {
                finalUrlString = urlModel.second;
            }
        }
    } else {
        // It isn't a default url, but an url from treeWidget
        QList<QTreeWidgetItem *> resultList_ = m_ui->m_bookmark->m_vulnBookmarkWidget->treeWidgetVulnUrl->findItems(m_vulnerabilityWidget->comboWebV->currentText(), Qt::MatchExactly, 0);
        finalUrlString = resultList_[0]->text(1);
    }

    finalUrlString.append(tmpSearchString);

    QUrl urlFinal(finalUrlString);
    searchVulnerability(urlFinal, m_vulnerabilityWidget->comboVulnRis->lineEdit()->text());
}

void Vulnerability::searchVulnerability(const QUrl& address, const QString& tabName)
{
    // default action settings
    m_collections.value("stop-act")->setEnabled(true);
    m_collections.value("back-act")->setEnabled(true);
    m_collections.value("forward-act")->setEnabled(true);

    m_welcomeUi->setVisible(false);
    // make tabBar visible for the first search
    if (!m_webViewList.size()) {
        m_vulnerabilityWidget->tWresult->setVisible(true);
    }

    // create new tab with a QWebView
    QWebView *page = new QWebView(m_ui);
    m_webViewList.push_back(page);
    page->load(address);
    m_vulnerabilityWidget->tWresult->addTab(page, tabName);
    connect(page, SIGNAL(loadProgress(int)), m_vulnerabilityWidget->progressWeb, SLOT(setValue(int)));
    connect(page, SIGNAL(loadFinished(bool)), this, SLOT(vulnPostScan()));
}

void Vulnerability::searchVulnFromTreeWidget()
{
    if (!m_vulnerabilityWidget->treeVulnNseRecovered->selectedItems().size()) {
        return;
    }

    QString addressName = m_ui->m_scanWidget->treeMain->selectedItems()[0]->text(0).split('\n', QString::SkipEmptyParts)[0];

    searchVulnerability(m_vulnerabilityWidget->treeVulnNseRecovered->selectedItems()[0]->text(0),
                        addressName
                        + "::"
                        + static_cast<QUrl>(m_vulnerabilityWidget->treeVulnNseRecovered->selectedItems()[0]->text(0)).authority());
}

void Vulnerability::closeVulnTab(int index)
{
    QWebView *page = m_webViewList.takeAt(index);
    m_vulnerabilityWidget->tWresult->removeTab(index);
    delete page;

    if (!m_webViewList.size()) {
        m_vulnerabilityWidget->tWresult->setVisible(false);
        m_welcomeUi->setVisible(true);
        // disable search action in vulnerability toolBar
        m_collections.value("search-act")->setEnabled(false);
    }
}

void Vulnerability::tabWebBack()
{
    m_webViewList[m_vulnerabilityWidget->tWresult->currentIndex()]->triggerPageAction(QWebPage::Back);
}

void Vulnerability::tabWebForward()
{
    m_webViewList[m_vulnerabilityWidget->tWresult->currentIndex()]->triggerPageAction(QWebPage::Forward);
}

void Vulnerability::tabWebStop()
{
    m_webViewList[m_vulnerabilityWidget->tWresult->currentIndex()]->triggerPageAction(QWebPage::Stop);
}

void Vulnerability::vulnPostScan()
{
    m_collections.value("stop-act")->setEnabled(false);
}

void Vulnerability::callVulnCheck()
{
    m_vulnerabilityWidget->comboVulnRis->clear();
    m_vulnerabilityWidget->comboVulnRis->insertItem(0, m_ui->m_bookmark->m_vulnBookmarkWidget->treeBookVuln->currentItem()->text(0));
    searchVulnerabilityFromCombo();
}

void Vulnerability::linkCompleterToServices()
{
    if (!m_ui->actionAdd_Bookmark->isEnabled()) {
        m_ui->actionAdd_Bookmark->setEnabled(true);
    }


    if (m_vulnModel.isNull()) {
        return;
    }

    if (m_completerVuln.isNull()) {
        m_completerVuln = new QCompleter(m_vulnModel.data(), m_ui);
        m_completerVuln.data()->setCompletionRole(QCompleter::InlineCompletion);
        m_completerVuln.data()->setWrapAround(false);
        m_completerVuln.data()->setCaseSensitivity(Qt::CaseInsensitive);
        m_vulnerabilityWidget->comboVulnRis->setCompleter(m_completerVuln.data());
    }
}

void Vulnerability::updateComboVuln(const QString& value)
{
    m_vulnerabilityWidget->comboVulnRis->clear();

    if (m_vulnerabilityWidget->comboVuln->currentIndex()) {
        m_vulnerabilityWidget->comboVulnRis->insertItem(0, value);
    } else {
        m_vulnerabilityWidget->comboVulnRis->lineEdit()->clear();
    }

}

void Vulnerability::objVulnButton()
{
    // SLOT
    m_vulnerabilityWidget->comboVulnRis->clear();

    QString tmpServices_ = m_ui->m_scanWidget->GItree->currentItem()->text(0);
    QList<QTreeWidgetItem*> itemList = m_ui->m_scanWidget->listWscan->findItems(tmpServices_, Qt::MatchExactly, 2);

    if (m_vulnerabilityWidget->comboVulnRis->itemText(0).isEmpty()) {
        m_vulnerabilityWidget->comboVulnRis->addItem(itemList[0]->text(3));
    } else {
        m_vulnerabilityWidget->comboVulnRis->setItemText(0, itemList[0]->text(3));
    }

    m_ui->m_collections->m_collectionsButton.value("scan-sez")->setChecked(false);
    m_ui->m_collections->m_collectionsButton.value("vuln-sez")->setChecked(true);
    m_ui->m_collections->m_collectionsButton.value("discover-sez")->setChecked(false);

    m_ui->toolBar->setVisible(false);
    m_ui->toolBarBook->setVisible(true);
    m_ui->toolBar_2->setVisible(false);
    m_ui->toolBarSearch->setVisible(true);

    m_ui->toolMenuBar->setVisible(false);
    m_ui->m_collections->m_discoverToolBar->setVisible(false);

    m_ui->m_mainTabWidget->removeTab(m_ui->m_mainTabWidget->indexOf(m_ui->m_bookmark->m_scanBookmarkWidget));
    m_ui->m_mainTabWidget->removeTab(m_ui->m_mainTabWidget->indexOf(m_ui->m_scanWidget));
    m_ui->m_mainTabWidget->removeTab(m_ui->m_mainTabWidget->indexOf(m_ui->m_discoverManager->m_discoverWidget));
    m_ui->m_mainTabWidget->removeTab(m_ui->m_mainTabWidget->indexOf(m_ui->m_bookmark->m_vulnBookmarkWidget));
    m_ui->m_mainTabWidget->insertTab(0, m_vulnerabilityWidget, QIcon(QString::fromUtf8(":/images/images/viewmag+.png")), "Vulnerability");
    m_ui->m_mainTabWidget->setCurrentIndex(0);

    // disable scan action
    m_ui->m_collections->disableScanUiActions();

    m_ui->menu_Bookmaks->setEnabled(true);
    m_ui->actionAdd_Bookmark->setText(tr("Add service to &bookmark"));
    m_ui->actionAdd_Parameters_to_bookmark->setVisible(false);
    m_ui->actionAdd_vulnerability_search_url->setVisible(true);
}

void Vulnerability::loadSearchUrlIntoCombo()
{
    // clear comboWeb
    m_vulnerabilityWidget->comboWebV->clear();

    // insert default static url in comboWeb
    QListIterator< QPair<QString, QString> > i(getDefaultUrlList());
    while (i.hasNext()) {
        m_vulnerabilityWidget->comboWebV->insertItem(m_vulnerabilityWidget->comboWebV->count() + 1, i.next().first);
    }

    m_vulnerabilityWidget->comboWebV->insertSeparator(m_vulnerabilityWidget->comboWebV->count() + 1);

    // value from treeWidget url
    for (int index = 0; index < m_ui->m_bookmark->m_vulnBookmarkWidget->treeWidgetVulnUrl->topLevelItemCount(); index++) {
        m_vulnerabilityWidget->comboWebV->insertItem(m_vulnerabilityWidget->comboWebV->count() + 1,
                m_ui->m_bookmark->m_vulnBookmarkWidget->treeWidgetVulnUrl->topLevelItem(index)->text(0));
    }
}

void Vulnerability::showAddUrlUi()
{
    QWeakPointer<addVulnerabilityUrl> dialogUrl = new addVulnerabilityUrl(m_ui);

    connect(dialogUrl.data(), SIGNAL(doneUrl(QString,QString)),
            this, SLOT(addUrlToBookmarks(QString,QString)));

    dialogUrl.data()->exec();

    if (!dialogUrl.isNull()) {
        delete dialogUrl.data();
    }
}

void Vulnerability::addUrlToBookmarks(const QString urlName, const QString urlAddr)
{
    m_ui->m_bookmark->saveAddressToBookmark(urlName, urlAddr);
    // reload comboVuln with website
    loadSearchUrlIntoCombo();
}

void Vulnerability::removeUrlToBookmarks()
{
    if (m_ui->m_bookmark->m_vulnBookmarkWidget->treeWidgetVulnUrl->selectedItems().isEmpty()) {
        return;
    }

    //remove address from bookmark
    m_ui->m_bookmark->deleteAddressFromBookmark(m_ui->m_bookmark->m_vulnBookmarkWidget->treeWidgetVulnUrl->selectedItems()[0]->text(0));
    loadSearchUrlIntoCombo();
}

void Vulnerability::updateComboServicesProperties()
{
    m_vulnerabilityWidget->comboVulnRis->clear();
    m_collections.value("search-act")->setEnabled(true);
    m_vulnerabilityWidget->comboVulnRis->setStyleSheet(QString::fromUtf8(""));
    bool signalState = m_vulnerabilityWidget->comboVulnRis->lineEdit()->disconnect(SIGNAL(cursorPositionChanged(int,int)));

    if (!signalState)
        return;

    connect(m_vulnerabilityWidget->comboVulnRis, SIGNAL(editTextChanged(QString)),
            this, SLOT(linkCompleterToServices()));
}
